# -*- coding: utf-8 -*-
"""
水库调度规程智能问答系统 - 简洁版界面
参考 DeepSeek 设计风格
"""

import sys
import os
from pathlib import Path
import streamlit as st

# 将项目根目录添加到路径
project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate

from D_db.core.db_manager import DBManager
from D_db.core.tokenizer import Tokenizer
from D_db.tools.prompt_utils import build_llm_context_from_qa

# 加载环境变量
load_dotenv()

# 配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

# 页面配置
st.set_page_config(
    page_title="水库调度问答",
    page_icon="💧",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# 简洁的CSS样式
st.markdown("""
<style>
    /* 隐藏默认元素 */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
    
    /* 整体布局 */
    .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
        max-width: 900px;
    }
    
    /* 标题样式 */
    .main-title {
        text-align: center;
        color: #1a1a1a;
        font-size: 1.8rem;
        font-weight: 600;
        margin-bottom: 2rem;
        padding: 1rem 0;
    }
    
    /* 消息容器 */
    .stChatMessage {
        background-color: transparent;
        padding: 1.2rem 0;
    }
    
    /* 用户消息 */
    [data-testid="stChatMessageContent"] {
        background-color: #f8f9fa;
        padding: 1rem 1.2rem;
        border-radius: 12px;
        border: 1px solid #e9ecef;
    }
    
    /* 助手消息 */
    .assistant-msg {
        background-color: #ffffff;
        padding: 1rem 1.2rem;
        border-radius: 12px;
        border: 1px solid #e9ecef;
        line-height: 1.6;
    }
    
    /* 输入框 */
    .stChatInputContainer {
        border-top: 1px solid #e9ecef;
        padding-top: 1rem;
    }
    
    /* 侧边栏 */
    [data-testid="stSidebar"] {
        background-color: #fafafa;
    }
    
    /* 对话列表项 */
    .conversation-item {
        padding: 0.8rem 1rem;
        margin: 0.3rem 0;
        border-radius: 8px;
        background-color: white;
        border: 1px solid #e9ecef;
        cursor: pointer;
        transition: all 0.2s;
        display: flex;
        justify-content: space-between;
        align-items: center;
    }
    
    .conversation-item:hover {
        background-color: #f8f9fa;
        border-color: #dee2e6;
    }
    
    .conversation-item.active {
        background-color: #e7f3ff;
        border-color: #90caf9;
    }
    
    .conversation-title {
        flex: 1;
        font-size: 0.9rem;
        color: #333;
        overflow: hidden;
        text-overflow: ellipsis;
        white-space: nowrap;
    }
    
    .delete-btn {
        color: #999;
        font-size: 0.9rem;
        padding: 0.2rem 0.5rem;
        margin-left: 0.5rem;
    }
    
    .delete-btn:hover {
        color: #ff5252;
    }
    
    /* 按钮样式 */
    .stButton button {
        width: 100%;
        border-radius: 8px;
        border: 1px solid #e9ecef;
        background-color: white;
        color: #333;
        padding: 0.5rem 1rem;
        font-size: 0.9rem;
        transition: all 0.2s;
    }
    
    .stButton button:hover {
        background-color: #f8f9fa;
        border-color: #dee2e6;
    }
    
    /* 展开器样式 */
    .streamlit-expanderHeader {
        font-size: 0.9rem;
        color: #666;
    }
    
    /* 滚动条 */
    ::-webkit-scrollbar {
        width: 6px;
        height: 6px;
    }
    
    ::-webkit-scrollbar-track {
        background: #f1f1f1;
    }
    
    ::-webkit-scrollbar-thumb {
        background: #c1c1c1;
        border-radius: 3px;
    }
    
    ::-webkit-scrollbar-thumb:hover {
        background: #a8a8a8;
    }
</style>
""", unsafe_allow_html=True)

class StreamingQueryTester:
    """查询处理器"""
    
    def __init__(self):
        self.db_manager = DBManager()
        self.tokenizer = Tokenizer()
        
        self.llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.7,
            streaming=True
        )
        
        self.answer_template = PromptTemplate(
            input_variables=["question", "context"],
            template="""你是一个专业的水库调度规程问答助手。请根据以下参考信息，回答用户的问题。

用户问题：
{question}

参考信息：
{context}

要求：
1. 基于参考信息提供准确、完整的回答
2. 如果参考信息不足以回答问题，请明确指出
3. 回答要简洁明了，直接针对问题给出答案
4. 使用专业、客观的语气

请提供你的回答：
"""
        )
        
        self.answer_chain = self.answer_template | self.llm
    
    def query_database(self, test_question: str, limit: int = 20):
        """查询数据库"""
        try:
            self.db_manager.connect()
            
            token_data = self.tokenizer.tokenize_question(test_question)
            tokens = token_data.get('tokens', [])
            category_id = token_data.get('category_id')
            
            if not tokens:
                return None, []
            
            results = self.db_manager.query_qa(tokens, limit, category_id)
            context_for_llm, qa_results = build_llm_context_from_qa(results)
            
            return context_for_llm, qa_results
            
        finally:
            self.db_manager.close()
    
    def stream_answer(self, question: str, context: str):
        """流式生成回答"""
        try:
            for chunk in self.answer_chain.stream({"question": question, "context": context}):
                if hasattr(chunk, 'content'):
                    yield chunk.content
        except Exception as e:
            yield f"生成回答时出错: {str(e)}"

def initialize_session():
    """初始化会话"""
    # 对话历史管理
    if "conversations" not in st.session_state:
        st.session_state.conversations = {}
        st.session_state.conversation_titles = {}
        st.session_state.current_conv_id = "conv_1"
        st.session_state.conversations["conv_1"] = []
        st.session_state.conversation_titles["conv_1"] = "新对话"
        st.session_state.conv_counter = 1
    
    # 当前对话的消息
    if "messages" not in st.session_state:
        st.session_state.messages = st.session_state.conversations.get(st.session_state.current_conv_id, [])
    
    if "tester" not in st.session_state:
        try:
            st.session_state.tester = StreamingQueryTester()
            st.session_state.ready = True
        except Exception as e:
            st.session_state.ready = False
            st.session_state.error = str(e)

def create_new_conversation():
    """创建新对话"""
    st.session_state.conv_counter += 1
    new_id = f"conv_{st.session_state.conv_counter}"
    st.session_state.conversations[new_id] = []
    st.session_state.conversation_titles[new_id] = "新对话"
    st.session_state.current_conv_id = new_id
    st.session_state.messages = []

def switch_conversation(conv_id):
    """切换对话"""
    # 保存当前对话
    if st.session_state.current_conv_id in st.session_state.conversations:
        st.session_state.conversations[st.session_state.current_conv_id] = st.session_state.messages
    
    # 切换到新对话
    st.session_state.current_conv_id = conv_id
    st.session_state.messages = st.session_state.conversations.get(conv_id, [])

def delete_conversation(conv_id):
    """删除对话"""
    if len(st.session_state.conversations) <= 1:
        return  # 至少保留一个对话
    
    if conv_id in st.session_state.conversations:
        del st.session_state.conversations[conv_id]
        del st.session_state.conversation_titles[conv_id]
        
        # 如果删除的是当前对话，切换到第一个对话
        if st.session_state.current_conv_id == conv_id:
            first_conv_id = list(st.session_state.conversations.keys())[0]
            switch_conversation(first_conv_id)

def get_conversation_title(messages):
    """根据消息生成对话标题"""
    if not messages:
        return "新对话"
    
    first_user_msg = next((msg["content"] for msg in messages if msg["role"] == "user"), None)
    if first_user_msg:
        # 取前20个字符作为标题
        return first_user_msg[:20] + ("..." if len(first_user_msg) > 20 else "")
    return "新对话"

def update_conversation_title(conv_id):
    """更新对话标题"""
    messages = st.session_state.conversations.get(conv_id, [])
    st.session_state.conversation_titles[conv_id] = get_conversation_title(messages)

def main():
    """主函数"""
    initialize_session()
    
    # 侧边栏
    with st.sidebar:
        st.markdown("### 💧 水库调度问答")
        
        # 新建对话按钮
        if st.button("➕ 新建对话", use_container_width=True, type="primary"):
            create_new_conversation()
            st.rerun()
        
        st.divider()
        
        # 对话历史列表
        st.markdown("##### 对话历史")
        
        # 获取所有对话，按ID倒序排列（最新的在上面）
        conv_ids = sorted(st.session_state.conversations.keys(), reverse=True)
        
        for conv_id in conv_ids:
            col1, col2 = st.columns([4, 1])
            
            with col1:
                # 对话标题按钮
                title = st.session_state.conversation_titles.get(conv_id, "新对话")
                is_active = conv_id == st.session_state.current_conv_id
                
                button_label = f"{'📝 ' if is_active else ''}{title}"
                
                if st.button(
                    button_label,
                    key=f"conv_{conv_id}",
                    use_container_width=True,
                    disabled=is_active
                ):
                    switch_conversation(conv_id)
                    st.rerun()
            
            with col2:
                # 删除按钮
                if len(st.session_state.conversations) > 1:
                    if st.button("🗑️", key=f"del_{conv_id}", help="删除对话"):
                        delete_conversation(conv_id)
                        st.rerun()
        
        st.divider()
        
        # 系统状态
        if st.session_state.ready:
            st.success("系统就绪", icon="✅")
        else:
            st.error("系统未就绪", icon="❌")
            st.caption(f"错误: {st.session_state.get('error', '未知')}")
        
        st.divider()
        
        # 查询设置
        st.markdown("##### 设置")
        limit = st.slider(
            "查询数量", 
            min_value=5, 
            max_value=50, 
            value=20, 
            step=5,
            label_visibility="collapsed",
            help="调整返回的知识库结果数量"
        )
        st.session_state.query_limit = limit
        
        st.divider()
        
        # 示例问题
        st.markdown("##### 示例问题")
        examples = [
            "大伙房水库的总库容是多少？",
            "大伙房水库的正常高水位是多少？",
            "大伙房水库的防洪标准是什么？",
            "大伙房水库的主要功能有哪些？",
        ]
        
        for ex in examples:
            if st.button(ex, key=f"ex_{ex}", use_container_width=True):
                st.session_state.example = ex
    
    # 主界面
    st.markdown('<div class="main-title">水库调度规程智能问答系统</div>', unsafe_allow_html=True)
    
    if not st.session_state.ready:
        st.warning("系统未就绪，请检查配置", icon="⚠️")
        return
    
    # 显示历史消息
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])
            
            # 显示参考资料
            if msg["role"] == "assistant" and msg.get("sources"):
                with st.expander("📚 查看参考资料", expanded=False):
                    for i, src in enumerate(msg["sources"][:3], 1):
                        st.markdown(f"**{i}.** {src.get('question', '')}")
                        st.caption(src.get('answer', ''))
                        st.divider()
    
    # 处理示例问题
    if hasattr(st.session_state, 'example'):
        prompt = st.session_state.example
        del st.session_state.example
    else:
        prompt = st.chat_input("给水库调度助手发送消息...")
    
    # 处理输入
    if prompt:
        # 显示用户消息
        with st.chat_message("user"):
            st.markdown(prompt)
        
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # 更新对话标题（如果是新对话的第一条消息）
        if len(st.session_state.messages) == 1:
            update_conversation_title(st.session_state.current_conv_id)
        
        # 生成回答
        with st.chat_message("assistant"):
            with st.spinner("思考中..."):
                try:
                    limit = st.session_state.get('query_limit', 20)
                    context, sources = st.session_state.tester.query_database(prompt, limit=limit)
                    
                    if not context:
                        response = "抱歉，没有找到相关信息。"
                        st.markdown(response)
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": response,
                            "sources": []
                        })
                    else:
                        # 流式输出
                        placeholder = st.empty()
                        full_response = ""
                        
                        for chunk in st.session_state.tester.stream_answer(prompt, context):
                            full_response += chunk
                            placeholder.markdown(full_response + "▌")
                        
                        placeholder.markdown(full_response)
                        
                        # 保存消息
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": full_response,
                            "sources": sources
                        })
                        
                        # 保存到对话历史
                        st.session_state.conversations[st.session_state.current_conv_id] = st.session_state.messages
                        
                        # 显示参考资料
                        if sources:
                            with st.expander("📚 查看参考资料", expanded=False):
                                for i, src in enumerate(sources[:3], 1):
                                    st.markdown(f"**{i}.** {src.get('question', '')}")
                                    st.caption(src.get('answer', ''))
                                    st.divider()
                
                except Exception as e:
                    st.error(f"出错了: {str(e)}")

if __name__ == "__main__":
    main()

